package org.mybatis.generator.plugins;

import java.util.List;
import java.util.Set;

import org.mybatis.generator.api.IntrospectedColumn;
import org.mybatis.generator.api.IntrospectedTable;
import org.mybatis.generator.api.PluginAdapter;
import org.mybatis.generator.api.dom.java.FullyQualifiedJavaType;
import org.mybatis.generator.api.dom.java.Interface;
import org.mybatis.generator.api.dom.java.JavaVisibility;
import org.mybatis.generator.api.dom.java.Method;
import org.mybatis.generator.api.dom.java.Parameter;
import org.mybatis.generator.api.dom.java.TopLevelClass;
import org.mybatis.generator.api.dom.xml.Document;
import org.mybatis.generator.api.dom.xml.TextElement;
import org.mybatis.generator.api.dom.xml.XmlElement;
import org.mybatis.generator.codegen.mybatis3.ListUtilities;
import org.mybatis.generator.codegen.mybatis3.MyBatis3FormattingUtilities;
import org.mybatis.generator.config.Context;
import org.mybatis.generator.utils.JavaMapperGeneratorHelper;
import org.mybatis.generator.utils.XMLMapperGeneratorHelper;

/**
 * <b><code>SelectCountPlugin</code></b>
 * <p/>
 * Description: 查询记录总条数
 * <p/>
 * <b>Creation Time:</b> 2020/06/08.
 */
public class SelectCountPlugin extends PluginAdapter {

	private final static String SELECT_COUNT = "selectCount";

	@Override
	public boolean validate(List<String> warnings) {
		return true;
	}

	@Override
	public boolean clientGenerated(Interface interfaze, TopLevelClass topLevelClass,
			IntrospectedTable introspectedTable) {
		if (introspectedTable.getTargetRuntime() == IntrospectedTable.TargetRuntime.MYBATIS3) {
			clientSelectCountMethodGenerated(interfaze, introspectedTable, context);
		}
		return super.clientGenerated(interfaze, topLevelClass, introspectedTable);
	}

	@Override
	public boolean sqlMapDocumentGenerated(Document document, IntrospectedTable introspectedTable) {
		if (introspectedTable.getTargetRuntime() == IntrospectedTable.TargetRuntime.MYBATIS3) {
			sqlMapSelectCountElementGenerated(document, introspectedTable, context);
		}
		return super.sqlMapDocumentGenerated(document, introspectedTable);
	}

	private void sqlMapSelectCountElementGenerated(Document document, IntrospectedTable introspectedTable,
			Context context) {

		String parameterType;
		if (introspectedTable.getRules().generateRecordWithBLOBsClass()) {
			parameterType = introspectedTable.getRecordWithBLOBsType();
		} else {
			parameterType = introspectedTable.getBaseRecordType();
		}

		XmlElement selectCountXmlElement = XMLMapperGeneratorHelper.baseElementGenerator(
				XMLMapperGeneratorHelper.SELECT, SELECT_COUNT, parameterType, XMLMapperGeneratorHelper.RESULT_TYPE,
				"java.lang.Integer");

		StringBuilder sbr = new StringBuilder();
		sbr.append("select ");
		List<IntrospectedColumn> columns = introspectedTable.getPrimaryKeyColumns();
		if (columns.size() == 1) {
			sbr.append(String.format("count(%s) ", columns.get(0).getActualColumnName()));
		} else {
			sbr.append("count(1) ");
		}

		sbr.append("from " + introspectedTable.getFullyQualifiedTableNameAtRuntime());
		selectCountXmlElement.addElement(new TextElement(sbr.toString()));
		XmlElement whereElement = new XmlElement("where");
		selectCountXmlElement.addElement(whereElement);
		List<IntrospectedColumn> columnList = ListUtilities
				.removeIdentityAndGeneratedAlwaysColumns(introspectedTable.getAllColumns());
		boolean and = false;
		for (int i = 0; i < columnList.size(); i++) {
			sbr.setLength(0);
			if (and) {
				sbr.append("  and ");
			} else {
				and = true;
			}
			IntrospectedColumn introspectedColumn = columnList.get(i);
			sbr.append(MyBatis3FormattingUtilities.getAliasedEscapedColumnName(introspectedColumn));
			sbr.append(" = ");
			sbr.append(MyBatis3FormattingUtilities.getParameterClause(introspectedColumn));
			XmlElement ifElement = XMLMapperGeneratorHelper
					.baseIfJudgeElementGen(introspectedColumn.getJavaProperty(), sbr.toString(), false);
			whereElement.addElement(ifElement);
		}
		document.getRootElement().addElement(selectCountXmlElement);
	}

	private void clientSelectCountMethodGenerated(Interface interfaze, IntrospectedTable introspectedTable,
			Context context) {
		Set<FullyQualifiedJavaType> importedTypes = JavaMapperGeneratorHelper
				.importedBaseTypesGenerator(introspectedTable);
		Method selectCountMethod = JavaMapperGeneratorHelper.methodGenerator(SELECT_COUNT, JavaVisibility.DEFAULT,
				FullyQualifiedJavaType.getIntInstance(),
				new Parameter(introspectedTable.getRules().calculateAllFieldsClass(), "record"));

		context.getCommentGenerator().addGeneralMethodComment(selectCountMethod, introspectedTable);
		interfaze.addImportedTypes(importedTypes);
		interfaze.addMethod(selectCountMethod);
	}
}
